import pandas as pd
import matplotlib.pyplot as plt
from xgboost import XGBRegressor
import joblib
from utils.data_util import preprocessing
from utils.log import Logger
import matplotlib.ticker as mick

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

# 设置显示的最大行数和列数为None
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)

# 创建日志对象
logging_obj = Logger(root_path="../",log_name="predict")
logger = logging_obj.get_logger()

def init():
    df = preprocessing("../data/test.csv")
    return df

def feature_extract(df):
    logger.info("开始进行特征提取")

    # 1- 基于time特征得到月份和小时
    df["month"] = pd.to_datetime(df["time"]).dt.month
    df["hour"] = pd.to_datetime(df["time"]).dt.hour

    # 2- 获取昨天同一时刻的用电负荷
    # 2.1- 将数据的格式转成字典。time作为key，power_load作为value
    time_load_dict = df.set_index("time")["power_load"].to_dict()
    # print(type(time_load_dict))
    # print(time_load_dict)

    def get_yesterday_load(row,interval):
        # 2.2- 基于当前日期时间获得昨天同一时刻的日期时间
        # 注意：current_time、yesterday_time不是字符串类型，而是DateTime日期时间对象类型
        current_time = pd.to_datetime(row["time"])
        yesterday_time = current_time - pd.to_timedelta(interval)

        # 2.3- 通过时刻获得负荷
        yesterday_time_str = yesterday_time.strftime("%Y-%m-%d %H:%M:%S")
        return time_load_dict.get(yesterday_time_str)

    # axis=1表示对DataFrame中的数据按行进行逐行处理
    df["yesterday"] = df.apply(get_yesterday_load,axis=1,args=("1d",))

    # 3- 获取当前时刻前1个小时、2个小时、3个小时的用电负荷
    # 实现方式一
    # df["pre_1_hour"] = df.apply(get_yesterday_load,axis=1,args=("1h",))
    # df["pre_2_hour"] = df.apply(get_yesterday_load,axis=1,args=("2h",))
    # df["pre_3_hour"] = df.apply(get_yesterday_load,axis=1,args=("3h",))

    # 实现方式二
    hours = 3
    shift_list = [df["power_load"].shift(i) for i in range(1,hours+1)]
    shift_data = pd.concat(shift_list,axis=1)
    shift_data.columns = [f"pre_{i}_hour" for i in range(1,hours+1)]
    # print(shift_list)
    # print(shift_data)
    df = pd.concat([df,shift_data],axis=1)

    # print(df.head(30))

    logger.info("完成特征提取")

    return df


def predict_data(df):
    # 1- 拆分得到特征数据
    all_columns = list(df.columns)
    all_columns.remove("power_load")
    all_columns.remove("time")

    x = df[all_columns]

    # 2- 特征预处理：标准化
    # 不能按照下面的写法对未知数据的特征进行处理。只能使用模型训练过程中的StandardScaler对未知数据使用相同的标准进行处理
    # transformer = StandardScaler()
    # x = transformer.transform(x)

    transformer = joblib.load("../model/transformer.pkl")
    x = transformer.transform(x)

    # 3- 加载训练好的模型对未知数据进行预测
    model:XGBRegressor = joblib.load("../model/power_load.pkl")
    y_predict = model.predict(x)

    return y_predict

def show_result(y_predict,init_df):
    # 创建画布
    fig = plt.figure(figsize=(40,20),dpi=150)
    ax = fig.add_subplot()

    # 绘制预测和真实的曲线
    ax.plot(init_df["time"],y_predict,label="预测结果") # 预测
    ax.plot(init_df["time"],init_df["power_load"],label="真实结果") # 真实

    # x轴刻度尺美化
    # 横坐标时间若不处理太过密集，这里调大时间展示的间隔
    ax.xaxis.set_major_locator(mick.MultipleLocator(50))
    # 时间展示时旋转45度
    plt.xticks(rotation=45)

    # 其他属性设置
    plt.title("预测结果展示")
    plt.legend()

    plt.savefig("../model/predict.jpg")
    plt.show()


if __name__ == '__main__':
    # 1- 初始化：读取要预测的数据
    init_df = init()

    # 2- 特征工程
    df = feature_extract(init_df)

    # 3- 未知数据预测
    y_predict = predict_data(df)

    # 4- 预测结果展示：将预测结果和真实结果进行比较
    show_result(y_predict,init_df)



